/* ================================================================
*   Copyright (C) 2020 All rights reserved.
*
*   文件名称：xkvstore.h
*   创 建 者：xunmenglong
*   创建日期：2020年12月05日
*   描    述：一个基于磁盘的KV存储库
*
================================================================ */


#ifndef XKVSTORE_H
#define XKVSTORE_H

#include <deque>

#include "xhashmap.h"
#include "xproto.h"
#include "xio_util.h"

using namespace xutil;

#define XKV_DEFAULT_BUCKET_NUM 1024*1024
#define XKV_DEFAULT_BLOCK_SIZE 536870912 //512M

class xkvstore {
private:
    struct mem_pos_t {
        int file_idx;       // 索引文件的idx
        int offset;         // 在该文件的哪个offset
        int len;            // 该data的长度
    };

#pragma pack (4)
    struct data_head_t {
        uint64_t key;       // 对应的key值, key值不能为0, 0要用作结束判断
        bool is_delete : 1; // 是否已经被删除
        uint body_len : 31;  // 数据长度
    };
#pragma pack()

private:
    xhashmap<uint64_t, mem_pos_t> _mem_idx;     // 内存索引
    deque<char *> _blocks;                      // 用来存储数据block列表, 即文件mmap的指针列表
    char * _cur_block = NULL;                   // 当前用于写的block
    int _cur_block_idx = -1;                    // 当前的索引block id
    int _cur_offset = 0;                        // 当前分配的block分配到什么位置了
    bthread_mutex_t _w_lock;                    // 写锁

    string _file_prefix;                        // 数据文件的前缀
    uint64_t _load_max_key;                          // 最大的key值

public:
    xkvstore() {}
    ~xkvstore() {
        // 将文件给ummap掉
        for (size_t i=0; i<_blocks.size(); i++) {
            char * block_ptr = _blocks[i];
            munmap((void *)block_ptr, XKV_DEFAULT_BLOCK_SIZE);
        }
    }

    /**
     * 初始化kvstore
     *
     * @param[in] file_prefix   数据文件的前缀
     * @param[in] bucket_num    内存索引桶数量, 根据你要存储的数据量来确定
     *                          如果内存空间足够，建议大小为你数据量/10
     * @return 0 SUCC, other FAIL
     */
    int init(string file_dir, int bucket_num = XKV_DEFAULT_BUCKET_NUM) {
        x_mkdir(file_dir.c_str());
        _file_prefix = file_dir + "/block.";
        _load_max_key = 10;

        int ret = _mem_idx.init(bucket_num);
        if (ret != 0) {
            LOG(ERROR) << "mem_idx init fail";
            return -1;
        }

        if (bthread_mutex_init(&_w_lock, NULL)) {
            LOG(ERROR) << "init write lock fail";
            return -1;
        }

        ret = _load_data();
        if (ret != 0) {
            LOG(ERROR) << "load data fail";
            return -1;
        }
        return 0;
    }

    /**
     * 插入数据，如果key值已经存在, 则直接覆盖老value
     * @TODO 需要做好老value占据空间的回收工作
     *
     * @param[in] key: 要插入的key值
     * @param[in] value: 要插入的value指针
     * @param[in] value_len: 要插入的数据的长度
     * @param[in] is_sync: 是否立即同步到磁盘
     *
     * @return 0 SUCC, other FAIL
     */
    int set(uint64_t key, const char * value, int value_len, bool is_sync=false) {
        LockUtil l(&_w_lock);

        if (key == 0) {
            LOG(ERROR) << "key should not be 0";
            return -1;
        }
        if (_cur_block == NULL ||
                _cur_offset+value_len > XKV_DEFAULT_BLOCK_SIZE) {
            // 需要分配新的block了
            char * new_block = _alloc_block();
            if (new_block == NULL) {
                LOG(ERROR) << "alloc block fail";
                return -1;
            }
            _blocks.push_back(new_block);
            _cur_block = new_block;
            _cur_offset = 0;
            _cur_block_idx++;
        }

        // 将数据存入磁盘
        int head_len = sizeof(data_head_t);
        int total_len = head_len + value_len;
        int pure_data_offset = _cur_offset + head_len;

        data_head_t cur_data;
        cur_data.key = key;
        cur_data.body_len = value_len;
        cur_data.is_delete = false;

        memcpy(_cur_block+_cur_offset, &cur_data, head_len);
        memcpy(_cur_block+pure_data_offset, value, value_len);
        if (is_sync) {
            int ret = x_msync(_cur_block+_cur_offset, total_len, MS_SYNC);
            if (ret < 0) {
                LOG(ERROR) << "msync fail: " << errno;
                return -1;
            }
        }
        _cur_offset += total_len;

        // 如果key值是已经存在的，将该key值的磁盘body_len设置为-1
        mem_pos_t * old_pos = _mem_idx.get(key);
        if (old_pos != NULL) {
            char * old_data_head_ptr = _blocks[old_pos->file_idx]+old_pos->offset - sizeof(data_head_t);
            data_head_t * old_data_head = (data_head_t *)old_data_head_ptr;
            old_data_head->is_delete = true;

            if (is_sync) {
                int ret = x_msync(old_data_head_ptr, sizeof(data_head_t), MS_SYNC);
                if (ret < 0) {
                    LOG(ERROR) << "msync fail: " << errno;
                    return -1;
                }
            }
        }

        // 存储索引
        mem_pos_t pos;
        pos.file_idx = _cur_block_idx;
        pos.offset = pure_data_offset;      // 指针指向纯数据的位置, 方便读取
        pos.len = value_len;
        _mem_idx.set(key, pos);

        return 0;
    }

    /**
     * 根据key值获取value
     *
     * @param[in] key: 要读取的key值
     * @param[out] out_value: 读取的value值存储的指针
     * @param[out] value_len: 读取的value值的长度
     *
     * @return 0 SUCC, 1 key不存在, other FAIL
     */
    int get(uint64_t key, char * out_value, int * value_len) {
        mem_pos_t * pos = _mem_idx.get(key);
        if (pos == NULL) {
            // key不存在
            return 1;
        }

        memcpy(out_value, _blocks[pos->file_idx]+pos->offset, pos->len);
        *value_len = pos->len;
        return 0;
    }

    /**
     * 删除指定key的数据
     *
     * @param[in] key: 要删除的key
     *
     * @return 0 SUCC, 1 key不存在, other FAIL
     */
    int remove(uint64_t key, bool is_sync=true) {
        mem_pos_t * pos = _mem_idx.get(key);
        if (pos == NULL) {
            // key值不存在
            return 1;
        }

        // 将磁盘上的数据标记删除
        char * data_head_ptr = _blocks[pos->file_idx]+pos->offset - sizeof(data_head_t);
        data_head_t * data_head = (data_head_t *)data_head_ptr;
        data_head->is_delete = true;

        if (is_sync) {
            int ret = x_msync(data_head_ptr, sizeof(data_head_t), MS_SYNC);
            if (ret < 0) {
                LOG(ERROR) << "msync fail: " << errno;
                return -1;
            }
        }

        // 删除内存索引
        return _mem_idx.remove(key);
    }

    /**
     * 快速往文件后面快速插入数据，为了提高写速度
     * 不更新内存索引，也不加锁，不停的往文件后面写
     *
     * 仅用于离线快速build索引文件使用
     * 初始化的block文件列表要为空
     *
     * @param[in] key:   要存储的key
     * @param[in] value/value_len:   要存储的value
     *
     * @return 0 SUCC, other FAIL
     */
    int off_build(uint64_t key, char * value, int value_len) {
        // 离线build，直接不加锁了
        if (_cur_block == NULL ||
                _cur_offset+value_len > XKV_DEFAULT_BLOCK_SIZE) {
            // 需要分配新的block了
            // 第一步, 将当前block给 munmap 掉
            if (_cur_block != NULL) {
                munmap((void *)_cur_block, XKV_DEFAULT_BLOCK_SIZE);
            }
            // 第二步, 分配新的block
            char * new_block = _alloc_block();
            if (new_block == NULL) {
                LOG(ERROR) << "alloc block fail";
                return -1;
            }
            _cur_block = new_block;
            _cur_offset = 0;
            _cur_block_idx++;
        }

        // 将数据存入磁盘
        int head_len = sizeof(data_head_t);
        int total_len = head_len + value_len;
        int pure_data_offset = _cur_offset + head_len;

        data_head_t cur_data;
        cur_data.key = key;
        cur_data.body_len = value_len;
        cur_data.is_delete = false;

        memcpy(_cur_block+_cur_offset, &cur_data, head_len);
        memcpy(_cur_block+pure_data_offset, value, value_len);
        _cur_offset += total_len;

        return 0;
    }

    inline uint64_t get_load_max_key() {
        return _load_max_key;
    }

private:
    // 从磁盘上加载数据，建立内存索引
    int _load_data() {
        int cur_idx = 0;
        while (true) {
            string cur_file = _file_prefix + to_string(cur_idx);
            if (is_file_exists(cur_file.c_str())) {
                // 有索引文件, 开始加载
                char * cur_block = x_mmap_file(cur_file.c_str());
                if (cur_block == NULL) {
                    LOG(ERROR) << "mmap block: " << cur_file << " fail";
                    return -1;
                }
                int last_offset = 0;
                int ret = _load_block(cur_block, cur_idx, &last_offset);
                if (ret != 0) {
                    LOG(ERROR) << "load block: " << cur_file << " fail";
                    return -1;
                }
                // 将当前块加入列表
                _blocks.push_back(cur_block);
                LOG(NOTICE) << cur_file << " load over";
                bool is_last_file = !(is_file_exists((_file_prefix+to_string(cur_idx+1)).c_str()));
                if (is_last_file) {
                    _cur_block = cur_block;
                    _cur_block_idx = cur_idx;
                    _cur_offset = last_offset;
                    break;
                }
            } else {
                // 一个文件都没有
                return 0;
            }
            cur_idx++;
        }
        return 0;
    }

    // 加载单个块文件
    int _load_block(char * block, int cur_idx, int * out_offset) {
        int cur_offset = 0;
        while (true) {
            if (cur_offset >= XKV_DEFAULT_BLOCK_SIZE) {
                // 文件已经遍历完了, 可以退出了
                break;
            }
            data_head_t * dh = (data_head_t *)(block+cur_offset);
            if (dh->key == 0) {
                // 文件还没遍历完, 但是没数据了
                break;
            }
            // key值为0是特殊值, 代表还没设置
            mem_pos_t pos;
            pos.file_idx = cur_idx;
            pos.offset = cur_offset + sizeof(data_head_t);
            pos.len = dh->body_len;
            if (!dh->is_delete) {
                int ret = _mem_idx.set(dh->key, pos);
                if (ret < 0) {
                    LOG(ERROR) << "mem_idx set " << dh->key << " fail";
                    return -1;
                }
                if (dh->key > _load_max_key) {
                    _load_max_key = dh->key;
                }
            }
            cur_offset += (sizeof(data_head_t)+dh->body_len);
        }
        *out_offset = cur_offset;
        return 0;
    }

    // 分配新块，并返回快指针
    char * _alloc_block() {
        string out_file = _file_prefix + to_string(_cur_block_idx + 1);
        FILE * fp;
        fp = fopen(out_file.c_str(), "w");
        //int ret = truncate(out_file.c_str(), XKV_DEFAULT_BLOCK_SIZE);
        int ret = ftruncate(fileno(fp), XKV_DEFAULT_BLOCK_SIZE);
        if (ret != 0) {
            LOG(ERROR) << "truncate " << out_file << errno << ": " << " fail";
            return NULL;
        }

        return x_mmap_file(out_file.c_str());
    }
};

#endif
